# 使用pythorch 完成
import torch
from torchvision.datasets import MNIST
from torchvision.transforms import Compose,ToTensor,Normalize
from torch.utils.data import DataLoader
from torch.optim import Adam
import torch.nn.functional as F
import torch.nn as nn
import os
import numpy as np

BATCH_SIZE = 128
TEST_BATCH_SIZE = 1000

# 1.准备数据集合
def get_dataloader(train=True,batch_size = BATCH_SIZE):
    transform_fn=Compose([ToTensor(),
             Normalize(mean=(0.1307,),std=(0.308,))])  # 多个transform连用

    dataset = MNIST(root="./data",train=train,transform=transform_fn)
    data_loader = DataLoader(dataset,batch_size=BATCH_SIZE,shuffle=True)
    return data_loader

# for i in enumerate(data_loader):
#     print(i)


# 2.构建模型

class MnistMoudel(nn.Module):
    def __init__(self):
        super(MnistMoudel, self).__init__()
        self.fc1 = nn.Linear(1*28*28,28)
        self.fc2 = nn.Linear(28,10)

    def forward(self,input):
        """

        :param input: [batch_size,1,28,28]
        :return:
        """
        # input.view([-1,1*28*28])
        x = input.view([input.size(0),1*28*28])

        #2.进行全连接
        x = self.fc1(x)
        #3.激活函数处理
        x = F.relu(x)
        #4.输出层
        out = self.fc2(x)
        return F.log_softmax(out,dim-1)

model = MnistMoudel()


optimizer = Adam(model.parameters(),lr=0.001)

if os.path.exists(("./model/model.pkl")):
    model.load_state_dict(torch.load("./model/model.pkl"))
    optimizer.load_state_dict(torch.load("./model/optimizer.pkl"))
def train(epoch):
    """实现训练实例化"""

    data_loader = get_dataloader(batch_size=BATCH_SIZE)

    for idx,(input,traget) in enumerate(data_loader):
        output = model(input) # 调用模型得到预测值
        # print(output)

        loss = F.nll_loss(output,traget) # 得到损失
        # print(traget)
        # exit()
        optimizer.zero_grad()  #梯度置零
        loss.backward() # 反向传播
        optimizer.step()  # 梯度更新

        if idx % 10 == 0:
            print(epoch,idx,loss.item())

        if idx %100 ==0:
            torch.save(model.state_dict(),"./model/model.pkl")
            torch.save(optimizer.state_dict(),"./model/optimizer.pkl")

def test():
    loss_list = []
    acc_list = []
    test_dataloader = get_dataloader(train= False,batch_size=TEST_BATCH_SIZE)
    for idx,(input,traget) in enumerate(test_dataloader):
        with torch.no_grad():  # 不对计算做追踪
            output = model(input)
            cur_loss = F.nll_loss(output,traget)
            loss_list.append(cur_loss)
            # 计算准确率
            pred = output.max(dim = -1)[-1]
            cur_acc = pred.eq(traget).float().mean()
            acc_list.append(cur_acc)
    print("平均准确率，平均损失",np.mean(acc_list),np.mean(loss_list))



if __name__ == "__main__":
    # for i in range(3):  # 训练三轮
    #     train(i)

    # loader = get_dataloader(train=False)
    # for input,label in loader:
    #     print(label)
    #     break
    test()




